import express, { Request, Response } from 'express'
import {
  RunAgentInputSchema,
  RunAgentInput,
  EventType,
  Message,
} from '@ag-ui/core'
import { EventEncoder } from '@ag-ui/encoder'
import { OpenAI } from 'openai'
import { v4 as uuidv4 } from 'uuid'

const app = express()

app.use(express.json())

const OPENAI_API_KEY = '';

app.post('/awp', async (req: Request, res: Response) => {
  try {
    console.log('req.body', typeof req.body, req.body) // 解析并验证请求体
    const input = req.body

    console.log('input', input) // 设置 SSE headers

    res.setHeader('Content-Type', 'text/event-stream')
    res.setHeader('Cache-Control', 'no-cache')
    res.setHeader('Connection', 'keep-alive') // 创建事件 encoder

    const encoder = new EventEncoder() // 发送 started 事件

    const runStarted = {
      type: EventType.RUN_STARTED,
      threadId: input.threadId,
      runId: input.runId,
    }
    res.write(encoder.encode(runStarted))

    if (OPENAI_API_KEY) {
      // 初始化 OpenAI 客户端
      const client = new OpenAI({ apiKey: OPENAI_API_KEY })

      // 将 AG-UI 消息转换为 OpenAI 消息格式
      const openaiMessages = input.messages
        .filter((msg: Message) =>
          ['user', 'system', 'assistant'].includes(msg.role),
        )
        .map((msg: Message) => ({
          role: msg.role as 'user' | 'system' | 'assistant',
          content: msg.content || '',
        })) // 生成消息 ID

      console.log('openaiMessages', openaiMessages)

      const messageId = uuidv4() // 发送 ‘文本消息开始’ 事件

      const textMessageStart = {
        type: EventType.TEXT_MESSAGE_START,
        messageId,
        role: 'assistant',
      }
      res.write(encoder.encode(textMessageStart)) // 创建流式传输完成请求

      const stream = await client.chat.completions.create({
        model: 'gpt-3.5-turbo',
        messages: openaiMessages,
        stream: true,
      }) // 处理流并发送 ‘文本消息内容’ 事件

      for await (const chunk of stream) {
        if (chunk.choices[0]?.delta?.content) {
          const content = chunk.choices[0].delta.content
          const textMessageContent = {
            type: EventType.TEXT_MESSAGE_CONTENT,
            messageId,
            delta: content,
          }
          res.write(encoder.encode(textMessageContent))
        }
      }

      // 发送 ‘文本消息结束’ 事件
      const textMessageEnd = {
        type: EventType.TEXT_MESSAGE_END,
        messageId,
      }
      res.write(encoder.encode(textMessageEnd))
    } else {
      const stream = [{
        choices: [{
          delta: {
            content: 'Hello! How can I help you today?',
          },
        }],
      }]
      for await (const chunk of stream) {
        const messageId = uuidv4() // 发送 ‘文本消息开始’ 事件
      
        if (chunk.choices[0]?.delta?.content) {
          const content = chunk.choices[0].delta.content
          const textMessageContent = {
            type: EventType.TEXT_MESSAGE_CONTENT,
            messageId,
            delta: content,
          }
          res.write(encoder.encode(textMessageContent))
        }
      }
    }

    // 发送 finished 事件
    const runFinished = {
      type: EventType.RUN_FINISHED,
      threadId: input.threadId,
      runId: input.runId,
    }
    res.write(encoder.encode(runFinished)) // 结束响应

    res.end()
  } catch (error) {
    console.log('error', error)
    res.status(422).json({ error: (error as Error).message })
  }
})

app.use('/', express.static('public'))

app.listen(8000, () => {
  console.log('Server running on http://localhost:8000')
})
